from .tools import get_func_dic_and_func_description, logger
import json
from typing import Callable
from .messages import (
    build_ai_message,
    build_human_message,
    build_tool_message,
    HistoryLoader, build_system_message,
)

class ChatAI:
    def __init__(self, history_loader: HistoryLoader = None):
        """
        构造一个可以使用的LLM
        :param history_loader: 和这个类一样，需要继承重写方法使用，历史记录不由此类及其子类管理；
        """
        # 工具的描述信息，给大模型的
        self.tools: list = []
        # 可调用的工具字典，通过字典调用函数
        self.tool_calls: dict[str, Callable] = {}
        # 临时内存历史，主要为了辅助工具的使用，工具需要对话历史才能调用
        # 在一次对话中，表明是一答一复，实际上加上可能的工具调用有很多次
        self.history_message: list = []
        # 加载器加载历史
        self.is_load_history = False
        self.history_loader: HistoryLoader = history_loader
        self.bind_history_loader(history_loader)
        self.model_name, self._client = self._get_llm_client()

    def bind_history_loader(self, history_loader: HistoryLoader):
        self.history_loader = history_loader
        if self.history_loader is not None:
            # 创建一个字典来保存被装饰的方法
            decorated_methods = [
                getattr(self.history_loader, method_name)
                for method_name in dir(self.history_loader)
                if
                callable(getattr(self.history_loader, method_name)) and getattr(getattr(self.history_loader, method_name), 'is_tool', False)
            ]
            funcs_dict, func_descriptions = get_func_dic_and_func_description(decorated_methods)
            # TODO 可能会有工具重名的情况，需要小心
            self.tools.extend(func_descriptions)
            self.tool_calls.update(funcs_dict)
            self.is_load_history = True

    def bind_tools(self, tools_list):
        """
        给大模型绑定一个可以使用的工具集合。
        :param tools_list: 工具列表；
        :return: 没有返回值；
        """
        funcs_dict, func_descriptions = get_func_dic_and_func_description(tools_list)
        self.tools.extend(func_descriptions)
        self.tool_calls.update(funcs_dict)

    def _get_llm_client(self):
        """
        需要重写这个方法
        :return: 返回第一个参数为要使用的model name，第二个参数为client;
        """
        raise NotImplementedError

    def add_message_to_history(self, message: dict | list):
        if isinstance(message, list):
            for m in message:
                self.history_message.append(m)
            return
        self.history_message.append(message)

    def merge_temporary_messages_into_history_loader(self):
        self.history_loader.add_messages(self.history_message)

    def __is_tool_call(self, response) -> bool:
        return response.choices[0].finish_reason == 'tool_calls'

    def __call_tools(self, response) -> str:
        for tool_call in response.choices[0].message.tool_calls:
            arguments = tool_call.function.arguments
            func_name = tool_call.function.name
            func_result = self.tool_calls[func_name](**json.loads(arguments))
            try:
                func_result_str = str(func_result)
            except TypeError as e:
                func_result_str = f"此工具返回的结果不能够转化为字符串类型，返回错误，不可调用: func_name: {func_name}"
                logger.error(func_result_str)

            tool_message = build_tool_message(func_result_str, tool_call.id, func_name, arguments)
            ai_message = build_ai_message(func_result_str)
            # 添加历史消息的时候，必须tool在前ai在后，顺序不能替换
            self.add_message_to_history(tool_message)
            # self.add_message_to_history(ai_message)
        # TODO 需要改成system消息
        return self.send_message_with_tools(
            build_system_message("注意查看是否还需要调用工具，如果需要请继续调用，否则根据已有的回答问题。", True)
        )

    def __message_to_dict(self, message):
        """
        智谱大模型只能处理字符串字典，这里做一个处理
        :param message: 这个响应是符合OpenAI的message.content规范的
        :return: 返回一个工具调用的信息，这个是弃用大模型tools的必要信息
        """
        # c = json.dumps(message.__dict__)
        if message.content is None: message.content = ''
        return {
            'role': str(message.role),
            'content': message.content,
            'tool_calls': [
                {
                    'function': {
                        'name': tool_call.function.name,
                        'arguments': tool_call.function.arguments,
                    },
                    'type': tool_call.type,
                    'id': tool_call.id,
                }
                for tool_call in message.tool_calls
            ],
        }

    def send_message_with_tools(self, message: dict) -> str:
        """
        需要使用build_XXX_message函数创建消息
        :param message: 输入的提问
        :return: 返回大模型的回答
        """
        if self.tools is None or len(self.tools) == 0:
            logger.warning("想要使用工具，但目前工具数为 0 。")

        history = self.history_loader.get_history() if self.is_load_history and self.history_loader else []
        self.add_message_to_history(message)

        completion = self._client.chat.completions.create(
            model=self.model_name,  # 填写需要调用的模型编码
            messages=history + self.history_message,
            tools=self.tools,
        )
        # 递归调用工具
        # TODO 可以改成while循环
        if self.__is_tool_call(completion):
            self.add_message_to_history(self.__message_to_dict(completion.choices[0].message))
            return self.__call_tools(completion)
        else:
            response = build_ai_message(completion.choices[0].message.content)

        # 只有不调用工具的时候，stop时，将消息加入到历史中
        self.add_message_to_history(response)
        if self.is_load_history: self.merge_temporary_messages_into_history_loader()
        # 每次会话都需要清除临时历史记录
        self.history_message = []
        return response['content']

    def send_message(self, message: str) -> str:
        history = self.history_loader.get_history() if self.is_load_history else []
        self.add_message_to_history(build_human_message(message))

        completion = self._client.chat.completions.create(
            model=self.model_name,  # 填写需要调用的模型编码
            messages=history + self.history_message,
            response_format={"type": "json_object"},
        )
        response = build_ai_message(completion.choices[0].message.content)
        self.add_message_to_history(response)
        if self.is_load_history: self.merge_temporary_messages_into_history_loader()
        # 每次会话都需要清除临时历史记录
        self.history_message = []
        return response['content']


